
#include "../iface/LibSSH2Interface.h"
#include "../iface/debug.h"
#include "../utils/BaseConvert.h"
#include "SSHClient.h"

/** Initialises local data in the SSH client class.
 * 
 * @param ikeyprint The key fingerprint expected, in binary format.
 * 
 * @param ikeyprintlen The length of the key fingerprint in binary
 * format given.
 * 
 * @param irealservername The real name of the remote server (used for
 * instance when the system is telescoping). */
void SSHClient::initClient(const char *ikeyprint, int ikeyprintlen, const char *irealservername)
{
   if (ikeyprint)
   {
      keyprint = (char *)malloc(sizeof(char)*ikeyprintlen);
      memcpy(keyprint, ikeyprint, sizeof(char)*ikeyprintlen);
      keyprintlen = ikeyprintlen;
   }
   realservername = (irealservername) ? strdup(irealservername) : strdup("");
   if (isvalid)
   {
      int retval = doSSHSessionHandshake();
      if (retval != 0)
      {
         invalidate();
         return;
      } else
      {
         // must set blocking!
         libssh2_session_set_blocking(tsession, 1);
         libssh2_trace(tsession, 
                       LIBSSH2_TRACE_SOCKET |
                       LIBSSH2_TRACE_TRANS |
                       LIBSSH2_TRACE_KEX |
                       LIBSSH2_TRACE_AUTH |
                       LIBSSH2_TRACE_CONN |
                       LIBSSH2_TRACE_SCP |
                       LIBSSH2_TRACE_SFTP |
                       LIBSSH2_TRACE_ERROR |
                       LIBSSH2_TRACE_PUBLICKEY);
      }
   }
   mutex = new APRMutex(mp);
}

/** Constructor for the SSH client. Since all we need is to open a TCP
 * client connection and then hijack the functionality with libssh2,
 * we can keep this quite simple, as this connects immediately.
 * 
 * @param ihostname Remote host name to connect to.
 * 
 * @param iportno Remote TCP port number to connect to.
 * 
 * @param imp Memory pool from which to derived the memory pool for
 * this class.
 * 
 * @param ikeyprint The key fingerprint expected, in binary format.
 * 
 * @param ikeyprintlen The length of the key fingerprint in binary
 * format given.
 * 
 * @param irealservername The real name of the remote server (used for
 * instance when the system is telescoping).
 * 
 * @param irealserverport The real port number of the remote server
 * (used for instance when the system is telescoping).
 * 
 * @param ipendingdb The pending database to add any problem
 * fingerprint into. */
SSHClient::SSHClient(const char *ihostname, int iportno, apr_pool_t *imp, const char *ikeyprint, int ikeyprintlen,
   const char *irealservername, int irealserverport, PendingDB *ipendingdb) :
   TCPClient(ihostname, iportno, imp), keyprint(NULL), keyprintlen(0), realserverport(irealserverport),
   localconnserverport(iportno), pendingdb(ipendingdb), mutex(NULL)
{
   initClient(ikeyprint, ikeyprintlen, irealservername);
}

/** Destructor for the SSH client. */
SSHClient::~SSHClient()
{
   if (keyprint) free(keyprint);
   keyprint = NULL;
   if (realservername) free(realservername);
   realservername = NULL;
   if (mutex) delete mutex;
   mutex = NULL;
   if (tsession != NULL)
   {
      libssh2_session_disconnect(tsession, "kthxbye");
      libssh2_session_free(tsession);
   }
}

/** Do the handshake to obtain a session.
 * 
 * @return Non-zero on failure, zero on success. */
int SSHClient::doSSHSessionHandshake()
{
   if (isvalid == 0)
   {
      errprint("invalid session");
      return -1;
   }
   tsession = libssh2_session_init_ex(NULL, NULL, NULL, this);
   if (tsession == NULL)
   {
      errprint("could not initialise SSH session structure");
      return -1;
   }
   apr_status_t status;
   apr_os_sock_t berk_socket_des;
   status = apr_os_sock_get(&berk_socket_des, tsocket);
   CheckAPRError(status);
   int ssh_retval = -1;
   int retries = 3;
   // This handshake seems to fail randomly too much... retry it a few times.
   while ((ssh_retval != 0) && (retries--))
   {
      ssh_retval = libssh2_session_handshake(tsession, berk_socket_des);
      tCheckLibSSH2Error(ssh_retval);
   }
   if (ssh_retval)
   {
      errprint("failed to achieve handshake with server");
      return -1;
   }
   bool try_add_to_pending = false;
   bool fatal = false;
   const char *host_keyprint = libssh2_hostkey_hash(tsession, LIBSSH2_HOSTKEY_HASH_SHA1);
   char *host_keyprint_hex = BaseConvert::convertBytesToHexadecimal(host_keyprint, 20);//sha1=20*8
   char *expected_keyprint_hex = NULL;
   if (keyprint)
   {
      if (memcmp(host_keyprint, keyprint, ((keyprintlen < 20) ? keyprintlen : 20)) != 0)
      {
         // key fingerprint doesn't match - fatal
         errprint("key fingerprint check failed");
         char *expected_keyprint_hex = BaseConvert::convertBytesToHexadecimal(keyprint, keyprintlen);
         errprint("expected fingerprint %s, obtained fingerprint was %s", expected_keyprint_hex, host_keyprint_hex);
         errprint("adding to pending changes");
         try_add_to_pending = true;
         fatal = true;
      }
   } else
   {
      errprint("no fingerprint specified - this is a security hole!");
      try_add_to_pending = true;
   }
   // try adding a problem fingerprint as an update to the pending
   // database
   if ((try_add_to_pending) && (pendingdb))
      pendingdb->addNewHostKeyUpdateRequest(realservername, (expected_keyprint_hex) ? (expected_keyprint_hex) : "(none)",
         host_keyprint_hex);
   free(host_keyprint_hex);
   if (expected_keyprint_hex) free(expected_keyprint_hex);
   expected_keyprint_hex = NULL;
   if (fatal) return -1;
   return 0;
}

/** Uses a SecurityChain to authenticate an SSH session.
 * 
 * @param secchain The SeucrityChain to use to authenticate the
 * session.
 * 
 * @param username The username to use with the security chain.
 * 
 * @return Non-zero on failure, zero on success. */
int SSHClient::useSecurityChain(SecurityChain *secchain, const char *username)
{
   if (isvalid == 0)
   {
      errprint("invalid session");
      return -1;
   }
   SecurityToken *stoken = NULL;
   char *userauthlist = libssh2_userauth_list(tsession, username, strlen(username));
   if (strstr(userauthlist, "none") != NULL)
   {
      errprint("ssh %s:%d - no authentication needed?", realservername, realserverport);
      if (libssh2_userauth_authenticated(tsession))
         return 0;
      return 1;
   }
   if ((stoken = secchain->getNextPasswordToken()) != NULL)
   {
      if (strstr(userauthlist, "password") != NULL)
      {
         int retval = stoken->authenticateSSHSession(tsession, username);
         //tCheckLibSSH2Error(retval);
         elCheckLibSSH2Error(this);
         if (libssh2_userauth_authenticated(tsession))
            return 0;
      } else
         stoken->addUnsuccessfulAuthenticationMarking();
   }
   if ((stoken = secchain->getNextInteractiveToken()) != NULL)
   {
      if (strstr(userauthlist, "keyboard-interactive") != NULL)
      {
         int retval = stoken->authenticateSSHSession(tsession, username);
         //tCheckLibSSH2Error(retval);
         elCheckLibSSH2Error(this);
         if (libssh2_userauth_authenticated(tsession))
            return 0;
      } else
         stoken->addUnsuccessfulAuthenticationMarking();
   }
   if ((stoken = secchain->getNextKeyToken()) != NULL)
   {
      if (strstr(userauthlist, "publickey") != NULL)
      {
         int retval = stoken->authenticateSSHSession(tsession, username);
         //tCheckLibSSH2Error(retval);
         elCheckLibSSH2Error(this);
         if (libssh2_userauth_authenticated(tsession))
            return 0;
      } else
         stoken->addUnsuccessfulAuthenticationMarking();
   }
   return 1;
}
